package com.edata.bigdata.nebulasource

import com.edata.bigdata.util.EDataUtils
import com.vesoft.nebula.client.graph.data.ValueWrapper
import org.apache.log4j.Logger
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.sources.{BaseRelation, InsertableRelation, TableScan}
import org.apache.spark.sql.types.{StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SQLContext}
import org.json.JSONObject

import scala.collection.mutable.ArrayBuffer


class NebulaDataSource(ctx: SQLContext, params: Map[String, String], schema_o: StructType)
  extends BaseRelation with TableScan with InsertableRelation with Serializable {
  @transient lazy val LOGGER = Logger.getLogger(this.getClass)

  override def sqlContext: SQLContext = ctx

  override def schema: StructType = schema_o

  override def buildScan(): RDD[Row] = {
    val nebula_addr = params.get("graph_addr").get
    val nebula_user = params.get("user").get
    val nebula_password = params.get("password").get
    val nebula_enableSSL = params.get("graph_ssl").get
    val nebula_timeout = params.get("timeout").get
    val ngql = params.get("ngql").get
    val nebulaConnector = new NebulaConnector(nebula_addr,
      nebula_user,
      nebula_password,
      nebula_enableSSL.toBoolean,
      nebula_timeout.toInt)
    LOGGER.info(s"Executing NGQL:${ngql}")
    val rs = nebulaConnector.submit(s"${ngql}")
    assert(rs.isSucceeded, rs.getErrorMessage)
    val colNames = rs.getColumnNames
    val rowSize = rs.rowsSize()
    val values_array = new ArrayBuffer[Row]()
    for (i <- 0 to rowSize - 1) {
      val r = rs.rowValues(i)
      val value_array = new ArrayBuffer[String]()
      colNames.forEach(col => {
        val vw = r.get(col)
        val value = convertValueWrapperToStr(vw)
        value_array.append(value)
      })
      values_array.append(Row.fromSeq(value_array))
    }
    nebulaConnector.close()
    val rdd = sqlContext.sparkContext.parallelize(values_array)
    rdd
  }

  def convertValueWrapperToStr(vw: ValueWrapper): String = {
    var value = ""
    if (vw.isEmpty || vw.isNull) value = ""
    if (vw.isBoolean) value = vw.asBoolean().toString
    if (vw.isGeography) value = vw.asGeography().toString
    if (vw.isString) value = vw.asString()
    if (vw.isVertex) {
      val node = vw.asNode()
      val label = node.labels().get(0)
      val valueJson = new JSONObject()
      val prosMap = node.properties(label)
      val prosJson = new JSONObject()
      val vid_val = convertValueWrapperToStr(node.getId)
      prosJson.put("VID", vid_val)
      prosMap.forEach((k, v) => {
        prosJson.put(k, convertValueWrapperToStr(v))
      })
      valueJson.put(label, prosJson)
      value = valueJson.toString()
    }
    if(vw.isEdge){
      val edge = vw.asRelationship()
      val label= edge.edgeName()
      val valueJson = new JSONObject()
      val prosMap = edge.properties()
      val prosJson = new JSONObject()
      val sid_val = convertValueWrapperToStr(edge.srcId())
      val tid_val = convertValueWrapperToStr(edge.dstId())
      prosJson.put("SID",sid_val)
      prosJson.put("TID",tid_val)
      prosMap.forEach((k,v)=>{
        prosJson.put(k,convertValueWrapperToStr(v))
      })
      valueJson.put(label,prosJson)
      value = valueJson.toString()
    }
    value
  }

  override def insert(data: DataFrame, overwrite: Boolean): Unit = {
    val nebula_addr = params.get("graph_addr").get
    val nebula_user = params.get("user").get
    val nebula_password = params.get("password").get
    val nebula_enableSSL = params.get("graph_ssl").get
    val nebula_timeout = params.get("timeout").get
    val schema = params.get("schema").get
    val space = params.get("space").get
    val label_type = params.get("label_type").get
    val label = params.get("label").get
    val gid = params.get("gid").get
    val gidAsProp = params.get("gidAsProp").get

    data.foreachPartition((iter: Iterator[Row]) => {
      val value_array = new ArrayBuffer[String]()
      iter.foreach((row: Row) => {
        val v_array = new ArrayBuffer[String]()
        val gid_array = new ArrayBuffer[String]()
        val gids = gid.split("\\|")
        gids.foreach(g => {
          val c_t = g.split(":")
          gid_array.append(EDataUtils.convertVToSaveStrByDataType(EDataUtils.findColValOfRowByColName(row, c_t(0), c_t(1)), c_t(1)))
        })
        if (gidAsProp.toBoolean) {
          val col_types = schema.split(",")
          col_types.foreach(col_type => {
            val c_t = col_type.split(":")
            val value = EDataUtils.convertVToSaveStrByDataType(EDataUtils.findColValOfRowByColName(row, c_t(0), c_t(1)), c_t(1))
            v_array.append(value)
          })
        } else {
          val col_types = schema.split(",").filterNot(col_type => col_type.equals(gid))
          col_types.foreach(col_type => {
            val c_t = col_type.split(":")
            val value = EDataUtils.convertVToSaveStrByDataType(EDataUtils.findColValOfRowByColName(row, c_t(0), c_t(1)), c_t(1))
            v_array.append(value)
          })
        }
        val values = s"${gid_array.mkString("->")}:(${v_array.mkString(",")})"
        value_array.append(values)
      })
      if (!value_array.isEmpty) {
        val cols = schema.split(",").map(c_t => c_t.split(":")(0))
        val ngql = s"USE ${space}; INSERT ${label_type} ${label}(${cols.mkString(",")}) VALUES ${value_array.mkString(",")}"
        LOGGER.info(s"Executing NGQL:${ngql}")
        val nebulaConnector: NebulaConnector = new NebulaConnector(nebula_addr,
          nebula_user,
          nebula_password,
          nebula_enableSSL.toBoolean,
          nebula_timeout.toInt)
        try {
          val rs = nebulaConnector.submit(ngql)
          if (!rs.isSucceeded) {
            LOGGER.error(rs.getErrorMessage)
          }
        } catch {
          case e: Exception => e.printStackTrace()
        } finally {
          nebulaConnector.close()
        }
      }

    })
  }


}



